package com.fanshuai.tree.bst;

import com.fanshuai.queue.LinkedQueue;
import com.fanshuai.queue.Queue;
import com.fanshuai.stack.LinkedStack;
import com.fanshuai.stack.Stack;
import com.fanshuai.tree.VisitorK;

import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;

/**
 * 二叉搜索树
 * @param <K>
 * @param <V>
 */
public class BinaryTree<K extends Comparable<K>, V> implements BST<K, V>, Iterable<K> {
    class Node {
        public K key;
        public V value;
        public Node left;
        public Node right;

        public Node(K key, V value) {
            this.key = key;
            this.value = value;
            this.left = null;
            this.right = null;
        }
    }

    private Node root;
    private int size;

    public BinaryTree() {
        root = null;
        size = 0;
    }

    @Override
    public void put(K key, V value) {
        root = put(root, key, value);
    }

    private Node put(Node node, K key, V value) {
        if (null == node) {
            size++;
            return new Node(key, value);
        }

        if (key.compareTo(node.key) < 0) {
            node.left = put(node.left, key, value);
        } else if (key.compareTo(node.key) > 0) {
            node.right = put(node.right, key, value);
        } else {
            node.value = value;
        }

        return node;
    }

    @Override
    public V get(K key) {
        if (null == root) {
            return null;
        }

        Node node = root;
        while (node != null) {
            if (node.key.equals(key)) {
                return node.value;
            }

            if (key.compareTo(node.key) < 0) {
                node = node.left;
            } else {
                node = node.right;
            }
        }

        return null;
    }

    @Override
    public void delete(K key) {
        root = delete(root, key);
    }

    private Node delete(Node node, K key) {
        if (null == node) {
            return null;
        }

        if (key.compareTo(node.key) < 0) {
            node.left = delete(node.left, key);
        } else if (key.compareTo(node.key) > 0) {
            node.right = delete(node.right, key);
        } else {
            size--;
            if (node.left == null) { //左子树为空，摘掉node节点，返回右子树
                Node right = node.right;
                releaseNode(node);
                return right;
            } else if (node.right == null) { //左子树不为空，右子树为空，摘掉node节点，返回左子树
                Node left = node.left;
                releaseNode(node);
                return left;
            } else { //左子树，右子树均不为空，寻找后继节点，交换后继节点与当前节点数据，删除后继节点
                Node next = findMin(node.right);
                node.key = next.key;
                node.value = next.value;

                node.right = delete(node.right, next.key);
            }
        }

        return node;
    }

    private void releaseNode(Node node) {
        if (node != null) {
            if (node.left != null) {
                node.left = null;
            }
            if (node.right != null) {
                node.right = null;
            }
            node = null;
        }
    }

    private Node findMin(Node node) {
        Node t = node;
        while (t.left != null) {
            t = t.left;
        }

        return t;
    }

    private Node findMax(Node node) {
        Node t = node;
        while (t.right != null) {
            t = t.right;
        }

        return t;
    }

    private void swapValue(Node p, Node q) {
        K k = p.key;
        V v = p.value;

        p.key = q.key;
        p.value = q.value;

        q.key = k;
        q.value = v;
    }

    @Override
    public List<K> keys() {
        List<K> keys = new ArrayList<>();
        fillKeySet(root, keys);
        return keys;
    }

    private void fillKeySet(Node node, List<K> keySet) {
        if (null == node) {
            return;
        }

        fillKeySet(node.left, keySet);
        keySet.add(node.key);
        fillKeySet(node.right, keySet);
    }

    /**
     * 先序遍历
     * @param node
     * @param visitor
     */
    public void preOrder(Node node, VisitorK<K, V> visitor) {
        if (null == node) {
            return;
        }

        visitor.visit(node.key, node.value);
        preOrder(node.left, visitor);
        preOrder(node.right, visitor);
    }

    /**
     * 中序遍历
     * @param node
     * @param visitor
     */
    public void inOrder(Node node, VisitorK<K, V> visitor) {
        if (null == node) {
            return;
        }

        inOrder(node.left, visitor);
        visitor.visit(node.key, node.value);
        inOrder(node.right, visitor);
    }

    /**
     * 后续遍历
     * @param node
     * @param visitor
     */
    public void afterOrder(Node node, VisitorK<K, V> visitor) {
        if (null == node) {
            return;
        }

        afterOrder(node.left, visitor);
        afterOrder(node.right, visitor);
        visitor.visit(node.key, node.value);
    }

    /**
     * 先序遍历
     * @param visitor
     */
    public void preOrderNonReverse(VisitorK<K, V> visitor) {
        Node node = root;
        Stack<Node> stack = new LinkedStack<>();
        visitor.visit(node.key, node.value);
        stack.push(root);

        while (!stack.isEmpty()) {
            while (node.left != null) {
                node = node.left;
                visitor.visit(node.key, node.value);
                stack.push(node);
            }

            node = stack.pop();
            if (node.right != null) {
                node = node.right;
                visitor.visit(node.key, node.value);
                stack.push(node);
            }
        }
    }

    /**
     * 中序遍历
     * @param visitor
     */
    public void inOrderNonReverse(VisitorK<K, V> visitor) {
        Node node = root;
        Stack<Node> stack = new LinkedStack<>();
        stack.push(root);

        while (!stack.isEmpty()) {
            while (node != null && node.left != null) {
                node = node.left;
                stack.push(node);
            }

            node = stack.pop();
            visitor.visit(node.key, node.value);
            if (node.right != null) {
                node = node.right;
                stack.push(node);
            }
        }
    }

    /**
     * 后续遍历
     * @param visitor
     */
    public void postOrderNonReverse(VisitorK<K, V> visitor) {
        Stack<Node> stack = new LinkedStack<>();
        stack.push(root);
        Stack<Node> stack1 = new LinkedStack<>();

        while (!stack.isEmpty()) {
            Node node = stack.pop();
            stack1.push(node);
            if (node.left != null) {
                stack.push(node.left);
            }
            if (node.right != null) {
                stack.push(node.right);
            }
        }

        while (!stack1.isEmpty()) {
            Node node = stack1.pop();
            visitor.visit(node.key, node.value);
        }
    }

    public void levelOrder(VisitorK<K, V> visitorK) {
        Queue<Node> queue = new LinkedQueue<>();
        queue.enqueue(root);

        while (!queue.isEmpty()) {
            Node node = queue.dequeue();
            visitorK.visit(node.key, node.value);

            if (node.left != null) {
                queue.enqueue(node.left);
            }
            if (node.right != null) {
                queue.enqueue(node.right);
            }
        }
    }

    @Override
    public int size() {
        return size;
    }

    @Override
    public boolean empty() {
        return size == 0;
    }

    @Override
    public Iterator<K> iterator() {
        Stack<Node> stack = new LinkedStack<>();
        stack.push(root);

        return new Iterator<K>() {
            Node node = root;

            @Override
            public boolean hasNext() {
                return !stack.isEmpty();
            }

            @Override
            public K next() {
                while (node != null && node.left != null) {
                    node = node.left;
                    stack.push(node);
                }

                node = stack.pop();
                K key = node.key;

                if (node.right != null) {
                    node = node.right;
                    stack.push(node);
                }

                return key;
            }
        };
    }
}
